#!/usr/bin/python
# encoding: utf-8

import sys
import os
import struct                   # to parse lines of data from files
from cStringIO import StringIO  # to handle transposing large files
from time import strftime       # to provide datestamps for transposed files
import numpy                    # needed to hold array for output and do maths
import numpy.ma                 # needed to hold array for output and do maths

## simple download progress ticker for use with hook in url retrieve
def progressTicker(count, blockSize, totalSize):
    
    nBlocks = totalSize / blockSize 
    stepSize = max(1, nBlocks / 20)  # deliberately int division, proofed against zero
    if(count % stepSize == 0):
        sys.stdout.write('.')
        sys.stdout.flush()


## three workhorse transpose data functions
## - one to transpose
## - one to create a list of grid ref indices
## - one to read in a transposed file

def transposeGrim(grim, out, verbose=0):
    
    # converts grim file from cell indexed to date indexed format
    # - growing strings by location in memory is abominably slow
    #   because strings are immutable and grow by copying into new memory
    #   so uses cStringIO to store the growing strings in pseudo files
    
    # grim='/Users/dorme/Desktop/tyn_sc_2.03/data/iavar.1901-2000.cld'
    # verbose=1000
    
    # open file
    try:
        dataFile = open(grim, 'r')
    except IOError:
        sys.stderr.write('Cannot open grim file: ' + grim + '\n')
        sys.exit()
    
    # read header
    header = [''] * 5
    for hdr in range(len(header)):
        header[hdr] = dataFile.readline()
    
    # parse out the number of rows per grid-ref - number of years.
    # hard coded location
    years = header[4].split('] [')[1]
    years = years.split('=')[1]
    years = years.split('-')
    startYear = int(years[0])
    endYear = int(years[1])
    nYear = endYear-startYear+1
    
    # intialise list of cStrings to hold strings
    # - need to create a null list and fill in otherwise
    #   all the locations point to the _same_ StringIO instance,
    #   which is a bit of a bugger.
    
    store =  [[ ] * 12 ] * nYear
    store =  [[StringIO() for m in range(12)] for y in range(nYear)]
    
    cellsProc = 0
    
    # iterate over the remaining lines of the file
    for line in dataFile:
        
        # test to see if we've hit a new cell and should therefore 
        # go around the years again 
        if line[0:8] == 'Grid-ref':
            currYear = 0
            cellsProc += 1
            if(verbose > 0):
                if (cellsProc % verbose) == 0 :
                    print 'Processing cell ' + str(cellsProc)
        else:
            # otherwise get the data out
            line = line.rstrip()
            line = struct.unpack('5s' * 12, line)
            for m in range(12):
                store[currYear][m].write(line[m])
            currYear += 1
    
    dataFile.close()
    
    # export
    try:
        outputFile = open(out, 'w')
    except IOError:
        sys.stderr.write('Cannot open ' + out + '\n')
        # sys.exit()
    
    outputFile.writelines('Transposed date structured version of ' + os.path.basename(grim) + ': original header follows.\n')
    outputFile.writelines('Created ' + strftime("%Y-%m-%d %H:%M:%S") + ' using tynsc20Tools written by David Orme.\n' )
    outputFile.writelines(header)
    
    for year in range(nYear):
        if(verbose > 1):
            print 'Writing year ' + str(year + 1) + ' of ' + str(nYear)
        for month in range(12):
            outputFile.write('Time-ref=' + str(month + 1) + ',' + str(startYear + year) + '\n')
            currMap = store[year][month].getvalue()
            outputFile.write(currMap + '\n')
    
    # close output file and the StringIO
    outputFile.close()
    
    for y in range(nYear):
        for m in range(12):
            store[y][m].close()


def getGridRefIndexGrim(grim, out):
    
    # returns an ordered file of Grid-ref indices 
    # grim = 'data/obs.clim6190.globe.lan.cld'
    # out  = 'transpose/transposedGridRefIndex.txt'
    # getIndexGrim('data/obs.clim6190.globe.lan.cld', 'transpose/transposedGridRefIndex.txt')
    
    # open input file
    try:
        dataFile = open(grim, 'r')
    except IOError:
        sys.stderr.write('Cannot open grim file: ' + grim + '\n')
        sys.exit()
    
    # read header
    header = [''] * 5
    for hdr in range(len(header)):
        header[hdr] = dataFile.readline()
    
    gridref=[]
    
    # iterate over the remaining lines of the file
    for line in dataFile:
        
        # test to see if we've hit a new cell 
        if line[0:8] == 'Grid-ref':
            gridref.append(line.rstrip()[9:])
    
    dataFile.close()
    
    # open ouput file
    try:
        outFile = open(out, 'w')
    except IOError:
        sys.stderr.write('Cannot open output file: ' + out + '\n')
        # sys.exit()    
    
    # dump them out
    [outFile.write(x + '\n') for x in gridref]
    outFile.close()


def readTransposeGrim(trgrim, dateKeys='', ticker=False):
    
    # reads all, unless a list of 'mm,yyyy' style keys is passed to dateKeys
    # trgrim = 'transpose/tr.hd.cgcm2.a1fi.2080s.cld'
    
    # open file
    try:
        dataFile = open(trgrim, 'r')
    except IOError:
        sys.stderr.write('Cannot open transposed grim file: ' + trgrim + '\n')
        sys.exit()
    
    # read header
    header = [''] * 7
    for hdr in range(len(header)):
        header[hdr] = dataFile.readline()
    
    # is it a transposed grim file
    if header[0][0:10] != 'Transposed':
        sys.stderr.write('Not a transposed grim file: ' + trgrim + '\n')
        sys.exit()
    
    dataDict = {}
    nFound = 0
    allFound = False
    # load the data
    while True:
        
        moyrline = dataFile.readline()
        dataline = dataFile.readline()
        
        # check for end of file
        if not moyrline or allFound: break
         
        moyrline = moyrline.rstrip()
        moyr = moyrline.split('=')[1]
        
        # append an entry to the data dictionary
        # - if needed, check against keys if not empty
        if dateKeys == '' or (moyr in dateKeys):
            
            # turn the long string (fixed width 5 character blocks)
            # from the data line into numpy arrays - easier  and faster manipulation
            # - this may be memory and time hungry
            dataline = dataline.rstrip()
            dataline = map(float, struct.unpack('5s' * 67420, dataline))
            dataline = numpy.array(dataline)
            dataDict[moyr] = dataline
            
            nFound += 1
            if not dateKeys == '' and nFound == len(dateKeys): break
            
            if ticker:            
                sys.stdout.write('.')
                if (nFound % 40) == 0: sys.stdout.write('\n')
                sys.stdout.flush()
            
    # end ticker
    if ticker:sys.stdout.write('\n')
    
    dataFile.close()
    
    return dataDict 


# read in annual time series for Tyndall SC 2.03
def readAnn(ann):
    
    # open file
    try:
        dataFile = open(ann, 'r')
    except IOError:
        sys.stderr.write('Cannot open ann file: ' + ann + '\n')
        sys.exit()
    
    # read header
    header = [''] * 4
    for hdr in range(len(header)):
        header[hdr] = dataFile.readline()
    
    # read the rest
    data = dataFile.readlines()
    dataFile.close()
    data = [x.rstrip() for x in data]
    data = [x.split() for x in data]
    
    # turn into a dictionary
    data = dict((int(k[0]), float(k[1])) for k in data)
    
    return data 


# old code - kept to cannibalize
def XXXXXwriteGrimToAsciiGrid(outbase, dictionary, gridindex):
    
    # reusable bits
    # embedded into main function - maybe cannibalise for a direct grim to ascii translator
    header = 'ncols        720\nnrows        360\nxllcorner    -180\nyllcorner    -90\ncellsize     0.5\nNODATA_value -9999\n'
    rowOrder = range(360)
    rowOrder.reverse()
    
    for eachkey in dictionary.keys():
        # set up a store prefilled with -9999
        ascii = numpy.zeros([360,720]) - 9999
        
        # fill in asciiList
        for c in range(len(gridindex)):
            row = gridindex[c][1] - 1
            col = gridindex[c][0] - 1
            ascii[row , col] = dictionary[eachkey][c]
        
        # write out asciiList as a grid
        outfile = outbase + eachkey
        
        try:
            out = open(outfile, 'w')
        except IOError:                                                  
            sys.stderr.write('Cannot open output file: ' + outfile + '\n')
            sys.exit()
        
        # write header to file
        out.write(header)
        
        # write rows to file
        for r in rowOrder:
            # space delimited strings
            rowVals = [' '+ str(round(v, 2)) for v in ascii[r,]]
            rowVals = ''.join(rowVals)
            # lose first space 
            rowVals = rowVals[1:]
            out.write(rowVals + '\n')
        
        out.close()
        



